#include "unetv1.h"
#include "stream_tools.h"
#include "deconv.h"
#include "im2col.h"
#include "conv.h"
#include "maxpool_2x2.h"

// Define the global weight_reshape variable
ap_uint<P_OCH * P_ICH * BIT_WGHT> weight_reshape[FOLD_O][FOLD_I][K*K];
ap_uint<CONV1_P_OCH * CONV1_P_ICH * BIT_WGHT> conv1_weight[CONV1_OCH/CONV1_P_OCH][CONV1_ICH/CONV1_P_ICH][K*K];
ap_uint<CONV2_P_OCH * CONV2_P_ICH * BIT_WGHT> conv2_weight[CONV2_OCH/CONV2_P_OCH][CONV2_ICH/CONV2_P_ICH][K*K];
ap_uint<CONV3_P_OCH * CONV3_P_ICH * BIT_WGHT> conv3_weight[CONV3_OCH/CONV3_P_OCH][CONV3_ICH/CONV3_P_ICH][K*K];
ap_uint<CONV4_P_OCH * CONV4_P_ICH * BIT_WGHT> conv4_weight[CONV4_OCH/CONV4_P_OCH][CONV4_ICH/CONV4_P_ICH][K*K];
ap_uint<CENTER_P_OCH * CENTER_P_ICH * BIT_WGHT> center_weight[CENTER_OCH/CENTER_P_OCH][CENTER_ICH/CENTER_P_ICH][K*K];
ap_uint<DECONV1_P_OCH * DECONV1_P_ICH * BIT_WGHT> deconv1_weight[DECONV1_OCH/DECONV1_P_OCH][DECONV1_ICH/DECONV1_P_ICH][K*K];
ap_uint<CONV5_P_OCH * CONV5_P_ICH * BIT_WGHT> conv5_weight[CONV5_OCH/CONV5_P_OCH][CONV5_ICH/CONV5_P_ICH][K*K];
ap_uint<DECONV2_P_OCH * DECONV2_P_ICH * BIT_WGHT> deconv2_weight[DECONV2_OCH/DECONV2_P_OCH][DECONV2_ICH/DECONV2_P_ICH][K*K];
ap_uint<CONV6_P_OCH * CONV6_P_ICH * BIT_WGHT> conv6_weight[CONV6_OCH/CONV6_P_OCH][CONV6_ICH/CONV6_P_ICH][K*K];
ap_uint<DECONV3_P_OCH * DECONV3_P_ICH * BIT_WGHT> deconv3_weight[DECONV3_OCH/DECONV3_P_OCH][DECONV3_ICH/DECONV3_P_ICH][K*K];
ap_uint<CONV7_P_OCH * CONV7_P_ICH * BIT_WGHT> conv7_weight[CONV7_OCH/CONV7_P_OCH][CONV7_ICH/CONV7_P_ICH][K*K];
ap_uint<DECONV4_P_OCH * DECONV4_P_ICH * BIT_WGHT> deconv4_weight[DECONV4_OCH/DECONV4_P_OCH][DECONV4_ICH/DECONV4_P_ICH][K*K];
ap_uint<CONV8_P_OCH * CONV8_P_ICH * BIT_WGHT> conv8_weight[CONV8_OCH/CONV8_P_OCH][CONV8_ICH/CONV8_P_ICH][K*K];

void unetv1(data_stream<BIT_ACTV>& in, axiu_stream<BIT_CONV>& out)
{
	#pragma HLS INTERFACE ap_ctrl_none port=return
	#pragma HLS INTERFACE axis register both port=in
	#pragma HLS INTERFACE axis register both port=out
	#pragma HLS DATAFLOW

    // Input expansion and first convolution layer
    data_stream<BIT_ACTV*CONV1_P_ICH> input_e("input_e");
    expandWidth<BIT_ACTV, CONV1_P_ICH * BIT_ACTV, CONV1_H * CONV1_W * CONV1_ICH>(in, input_e);

    data_stream<BIT_CONV*CONV1_P_OCH> input_copy1("input_copy1");
    data_stream<BIT_CONV*CONV1_P_OCH> input_copy2("input_copy2");
    copy_stream<BIT_ACTV*CONV1_P_ICH, CONV1_H * CONV1_W * CONV1_ICH>(input_e, input_copy1, input_copy2);

    data_stream<BIT_ACTV*CONV1_P_ICH> input_im2col("input_im2col");
    im2col<1, CONV1_H, CONV1_W, CONV1_ICH, K, 1, P, BIT_ACTV, CONV1_P_ICH, 0>(input_copy1, input_im2col);

    data_stream<BIT_CONV*CONV1_P_OCH> conv1("conv1");
    conv<CONV1_P_ICH, CONV1_P_OCH, CONV1_ICH, CONV1_OCH, K, BIT_ACTV, BIT_WGHT, BIT_CONV, CONV1_H * CONV1_W>(
        input_im2col, conv1, conv1_weight
    );

    data_stream<BIT_CONV*CONV1_P_OCH> conv1_copy1("conv1_copy1");
    data_stream<BIT_CONV*CONV1_P_OCH> conv1_copy2("conv1_copy2");
    copy_stream<BIT_CONV*CONV1_P_OCH, CONV1_H * CONV1_W>(conv1, conv1_copy1, conv1_copy2);

    // First max pooling
    data_stream<BIT_CONV*POOL1_P_CH> maxpool1("maxpool1");
    maxpool2x2<POOL1_P_CH, CONV1_OCH, CONV1_H, CONV1_W, 1, BIT_CONV>(conv1_copy1, maxpool1);

    // Second convolution layer
    data_stream<BIT_CONV*CONV2_P_ICH> maxpool1_im2col("maxpool1_im2col");
    im2col<1, CONV1_H/2, CONV1_W/2, CONV1_OCH, K, 1, P, BIT_CONV, CONV2_P_ICH, 0>(maxpool1, maxpool1_im2col);

    data_stream<BIT_CONV*CONV2_P_OCH> conv2("conv2");
    conv<CONV2_P_ICH, CONV2_P_OCH, CONV2_ICH, CONV2_OCH, K, BIT_CONV, BIT_WGHT, BIT_CONV, CONV2_H * CONV2_W>(
        maxpool1_im2col, conv2, conv2_weight
    );

    data_stream<BIT_CONV*CONV2_P_OCH> conv2_copy1("conv2_copy1");
    data_stream<BIT_CONV*CONV2_P_OCH> conv2_copy2("conv2_copy2");
    copy_stream<BIT_CONV*CONV2_P_OCH, CONV2_H * CONV2_W>(conv2, conv2_copy1, conv2_copy2);

    // Second max pooling
    data_stream<BIT_CONV*POOL2_P_CH> maxpool2("maxpool2");
    maxpool2x2<POOL2_P_CH, CONV2_OCH, CONV2_H, CONV2_W, 1, BIT_CONV>(conv2_copy1, maxpool2);

    // Third convolution layer
    data_stream<BIT_CONV*CONV3_P_ICH> maxpool2_im2col("maxpool2_im2col");
    im2col<1, CONV2_H/2, CONV2_W/2, CONV2_OCH, K, 1, P, BIT_CONV, CONV3_P_ICH, 0>(maxpool2, maxpool2_im2col);

    data_stream<BIT_CONV*CONV3_P_OCH> conv3("conv3");
    conv<CONV3_P_ICH, CONV3_P_OCH, CONV3_ICH, CONV3_OCH, K, BIT_CONV, BIT_WGHT, BIT_CONV, CONV3_H * CONV3_W>(
        maxpool2_im2col, conv3, conv3_weight
    );

    data_stream<BIT_CONV*CONV3_P_OCH> conv3_copy1("conv3_copy1");
    data_stream<BIT_CONV*CONV3_P_OCH> conv3_copy2("conv3_copy2");
    copy_stream<BIT_CONV*CONV3_P_OCH, CONV3_H * CONV3_W>(conv3, conv3_copy1, conv3_copy2);

    // Third max pooling
    data_stream<BIT_CONV*POOL3_P_CH> maxpool3("maxpool3");
    maxpool2x2<POOL3_P_CH, CONV3_OCH, CONV3_H, CONV3_W, 1, BIT_CONV>(conv3_copy1, maxpool3);

    // Fourth convolution layer
    data_stream<BIT_CONV*CONV4_P_ICH> maxpool3_im2col("maxpool3_im2col");
    im2col<1, CONV3_H/2, CONV3_W/2, CONV3_OCH, K, 1, P, BIT_CONV, CONV4_P_ICH, 0>(maxpool3, maxpool3_im2col);

    data_stream<BIT_CONV*CONV4_P_OCH> conv4("conv4");
    conv<CONV4_P_ICH, CONV4_P_OCH, CONV4_ICH, CONV4_OCH, K, BIT_CONV, BIT_WGHT, BIT_CONV, CONV4_H * CONV4_W>(
        maxpool3_im2col, conv4, conv4_weight
    );

        // Convert conv4 for concatenation - using copy_stream for skip connection
    data_stream<BIT_CONV*CONV4_P_OCH> conv4_copy1("conv4_copy1");
    data_stream<BIT_CONV*CONV4_P_OCH> conv4_copy2("conv4_copy2");
    copy_stream<BIT_CONV*CONV4_P_OCH, CONV4_H * CONV4_W>(conv4, conv4_copy1, conv4_copy2);

    // Fourth max pooling
    data_stream<BIT_CONV*POOL4_P_CH> maxpool4("maxpool4");
    maxpool2x2<POOL4_P_CH, CONV4_OCH, CONV4_H, CONV4_W, 1, BIT_CONV>(conv4_copy1, maxpool4);

    // Center convolution (bottom of U-Net)
    data_stream<BIT_CONV*CENTER_P_ICH> maxpool4_im2col("maxpool4_im2col");
    im2col<1, CONV4_H/2, CONV4_W/2, CONV4_OCH, K, 1, P, BIT_CONV, CENTER_P_ICH, 0>(maxpool4, maxpool4_im2col);

    data_stream<BIT_CONV*CENTER_P_OCH> center("center");
    conv<CENTER_P_ICH, CENTER_P_OCH, CENTER_ICH, CENTER_OCH, K, BIT_CONV, BIT_WGHT, BIT_CONV, CENTER_H * CENTER_W>(
        maxpool4_im2col, center, center_weight
    );

    // Upsampling path starts here
    // First upsampling: deconv center and concatenate with conv4
    data_stream<BIT_CONV*DECONV1_P_OCH> deconv1("deconv1");
    deconv<DECONV1_P_ICH, DECONV1_P_OCH, DECONV1_ICH, DECONV1_OCH, DECONV1_IH, DECONV1_IW, K, P, S, O_P, BIT_CONV, BIT_WGHT, BIT_CONV>(
        center, deconv1, deconv1_weight
    );

    // Concatenate skip connection
    data_stream<BIT_CONV> concat4("concat4");
    concat<BIT_CONV, CONV4_OCH, DECONV1_OCH, CONV5_H * CONV5_W>(conv4_copy2, deconv1, concat4);

    data_stream<BIT_CONV*CONV5_P_ICH> concat4_wide("concat4_wide");
    expandWidth<BIT_CONV, BIT_CONV*CONV5_P_ICH, CONV5_H * CONV5_W * CONV5_ICH>(concat4, concat4_wide);

    data_stream<BIT_CONV*CONV5_P_ICH> concat4_im2col("concat4_im2col");
    im2col<1, CONV5_H, CONV5_W, CONV5_ICH, K, 1, P, BIT_CONV, CONV5_P_ICH, 0>(concat4_wide, concat4_im2col);

    data_stream<BIT_CONV*CONV5_P_OCH> conv5("conv5");
    conv<CONV5_P_ICH, CONV5_P_OCH, CONV5_ICH, CONV5_OCH, K, BIT_CONV, BIT_WGHT, BIT_CONV, CONV5_H * CONV5_W>(
        concat4_im2col, conv5, conv5_weight
    );

    // Final output - for simplicity, just add last signal and output
    data_stream<BIT_CONV> s_output("s_output");
    reduceWidth<BIT_CONV*CONV5_P_OCH, BIT_CONV, CONV5_H * CONV5_W * CONV5_OCH>(conv5, s_output);

    add_last<BIT_CONV, CONV5_H * CONV5_W * CONV5_OCH>(s_output, out);
}
